package com.etc


import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.ALS
import org.apache.spark.sql.SparkSession

/**
  * ASL基于DataFrame的Demo
  */
object ASLdemo {


  case class Rating(userId: Int, itemId: Int, rating: Float)

  /**
    * 解析数据：将数据转换成Rating对象
    * @param str
    * @return
    */
  def parseRating(str: String): Rating = {
    val fields = str.split(",")
    assert(fields.size == 3)
    Rating(fields(0).toInt, fields(1).toInt, fields(2).toFloat)
  }


  def main(args: Array[String]): Unit = {
    //定义切入点
    val spark = SparkSession.builder().master("local").appName("ASL-DF-Demo").getOrCreate()
    //读取数据，生成RDD并转换成Rating对象
    import spark.implicits._
    val ratingsDF = spark.sparkContext.textFile("u.data").map(parseRating).toDF()
    //将数据随机分成训练数据和测试数据（权重分别为0.8和0.2）
    val Array(training, test) = ratingsDF.randomSplit(Array(0.8, 0.2))
    //定义ALS，参数初始化
    val als = new ALS().setRank(50)
      .setMaxIter(10)
      .setRegParam(0.01)
      .setUserCol("userId")
      .setItemCol("itemId")
      .setRatingCol("rating")
    //训练模型
    val model = als.fit(training)



    //推荐:每个用户推荐2个物品
    val r = model.recommendForAllUsers(2)

    //关闭冷启动（防止计算误差不产生NaN）
    model.setColdStartStrategy("drop")
    //预测测试数据
    val predictions = model.transform(test)

    //定义rmse误差计算器
    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")
      .setLabelCol("rating")
      .setPredictionCol("prediction")
    //计算误差
    val rmse = evaluator.evaluate(predictions)

    //打印训练数据
    training.foreach(x=>println("训练数据： "+x))
    //打印测试数据
    test.foreach(x=>println("测试数据： "+x))
    //打印推荐结果
    r.foreach(x=>print("用户 "+x(0)+" ,推荐物品 "+x(1)))
    //打印预测结果
    predictions.foreach(x=>print("预测结果:  "+x))
    //输出误差
    println(s"Root-mean-square error = $rmse")
  }

}
